use std::collections::{HashMap, HashSet};
use tracing::debug;

use bc_base::{TransactionId, transaction::Transaction};

use crate::transaction::TransactionConsensusError;

#[derive(Debug)]
pub(super) struct OrphanTransactions {
    orphan_txs: HashMap<TransactionId, Transaction>, // <tx-id, tx>
    orphan_txs_by_prev: HashMap<TransactionId, HashSet<TransactionId>>, // <prev-tx-id, set(tx-id1, tx-id2, ...)>
}

impl OrphanTransactions {
    pub fn new() -> Self {
        Self {
            orphan_txs: HashMap::new(),
            orphan_txs_by_prev: HashMap::new(),
        }
    }

    // void AddOrphanTx(const CDataStream& vMsg) {
    pub fn add_orphan_tx(
        &mut self,
        tx_id: &TransactionId,
        tx: Transaction,
    ) -> Result<(), TransactionConsensusError> {
        if self.orphan_txs.contains_key(tx_id) {
            return Err(TransactionConsensusError::OrphanTransactionAlreadyExists);
        }

        for tx_in in tx.get_inputs() {
            let prev_tx_id = tx_in.get_prev_out_point().get_tx_id();
            if !self.orphan_txs_by_prev.contains_key(prev_tx_id) {
                self.orphan_txs_by_prev.insert(*prev_tx_id, HashSet::new());
            }

            self.orphan_txs_by_prev
                .get_mut(prev_tx_id)
                .unwrap()
                .insert(*tx_id);
        }

        self.orphan_txs.insert(*tx_id, tx);

        Ok(())
    }

    // void EraseOrphanTx(uint256 hash) {
    pub fn remove_orphan_tx(
        &mut self,
        tx_id: &TransactionId,
    ) -> Result<(), TransactionConsensusError> {
        let Some(tx) = self.orphan_txs.get(tx_id) else {
            debug!("TX ({}) is not Orphan transaction.", tx_id);
            return Ok(());
        };

        for tx_in in tx.get_inputs() {
            let prev_tx_id = tx_in.get_prev_out_point().get_tx_id();
            if let Some(set) = self.orphan_txs_by_prev.get_mut(prev_tx_id) {
                set.remove(tx_id);
                if set.is_empty() {
                    self.orphan_txs_by_prev.remove(prev_tx_id);
                }
            }
        }

        self.orphan_txs.remove(tx_id);

        Ok(())
    }

    pub fn get_orphan_txs_by_prev(
        &self,
        prev_tx_id: &TransactionId,
    ) -> Option<&HashSet<TransactionId>> {
        self.orphan_txs_by_prev.get(prev_tx_id)
    }
}
